import numpy as np
from util.plot import set_legend, set_axis
import matplotlib.pyplot as plt
from decimal import Decimal

from cal_op import cal_op
from deal_txt import read_txt, save_txt
from deal_sgy import read_sgy
from draw_data import draw_sgy
from draw_vel import draw_vel, calLayerVel


def main(param):
    # 自动拾取， 使用Decimal将数字字符串转为数字，不影响其精度
    auto_op = read_txt(param["path_auto_op"])  # 自动拾取的曲线
    auto_time = [Decimal(s) for s in auto_op["time"]]
    auto_depth = [Decimal(s) for s in auto_op["zrp"]]
    # 手动拾取
    hand_time = []
    if param["path_hand_op"]:
        hand_op = read_txt(param["path_hand_op"])  # 手动拾取的曲线
        hand_time = [Decimal(s) for s in hand_op["time"]]
    # 程序计算
    cal_time = cal_op(auto_time, auto_depth, param["cal_depth_start"], param["cal_depth_end"])  # 计算初至位置

    # ============================================================
    # figure1, 绘制sgy数据
    # ============================================================

    fig_idx = 0  # fig编号

    if param["path_sgy"]:
        plt.figure(fig_idx)  # 图形编号
        fig_idx += 1

        sgy = read_sgy(param["path_sgy"])  # sgy数据
        ax_sgy, axis = draw_sgy(plt, sgy, param["sgy_depth_start"], param["sgy_depth_end"], param["sgy_time_start"], param["sgy_time_end"])  # 绘制sgy数据
        auto_depth2 = auto_depth[axis["top"]:axis["bottom"]]  # 根据要显示的数据范围对auto_depth进行截取

        legend_flag = 0  # 是否调用set_legend设置图例和样式

        # 绘制自动拾取的曲线
        if param["show_sgy_auto"]:
            legend_flag = 1
            auto_time2 = auto_time[axis["top"]:axis["bottom"]]  # 根据要显示的数据范围对auto_time进行截取
            ax_sgy.plot(auto_time2, auto_depth2, 'go-', lw=0.24, markersize=0.3, label='auto')  # 绘制初至
            # 绘制自动拾取的中值线
            for i in range(len(auto_depth2)):
                ax_sgy.plot([auto_time2[i], auto_time2[i]], [auto_depth2[i], auto_depth2[i] - 10], 'g', linewidth=0.24)

        # 绘制手动拾取的曲线
        if param["path_hand_op"] and param["show_sgy_hand"]:
            legend_flag = 1
            hand_time2 = hand_time[axis["top"]:axis["bottom"]]  # 根据要显示的数据范围对hand_time进行截取
            ax_sgy.plot(hand_time2, auto_depth2, 'ro-', lw=0.24, markersize=0.3, label='hand')

        # 绘制计算的曲线
        if param["show_sgy_cal"]:
            legend_flag = 1
            cal_time2 = cal_time[axis["top"]:axis["bottom"]]  # 根据要显示的数据范围对cal_time进行截取
            ax_sgy.plot(cal_time2, auto_depth2, 'bo-', lw=0.24, markersize=0.3, label='cal')
            # 绘制计算初至的中值线
            for i in range(len(auto_depth2)):
                ax_sgy.plot([cal_time2[i], cal_time2[i]], [auto_depth2[i], auto_depth2[i] - 10], 'b', linewidth=0.24)

        # 设置图例样式和位置
        if legend_flag:
            set_legend(plt)

    # ============================================================
    # figure2
    # ============================================================

    # 计算深度差
    depth_diff = []
    depth_flag = []  # 深度差变化的位置
    depth_i = 10
    for i in range(len(auto_depth)):
        if i == 0:
            depth_diff.append(0)
        else:
            diff = auto_depth[i] - auto_depth[i - 1]
            depth_diff.append(diff)
            if depth_i != diff:
                depth_i = diff
                depth_flag.append(auto_depth[i])
    depth_diff[0] = depth_diff[1]

    # ========== 层速度 ==========

    plt.figure(fig_idx)  # 图形编号
    ax_vel, fig_vel = draw_vel(plt, auto_time, auto_depth)

    # 手动拾取
    if param["path_hand_op"]:
        hand_vel = calLayerVel(np.array(hand_time), np.array(auto_depth))
        ax_vel.plot(hand_vel, np.array(auto_depth), 'r', lw=0.24, label='hand')

    # 计算初至
    cal_vel = calLayerVel(np.array(cal_time), np.array(auto_depth))
    ax_vel.plot(cal_vel, np.array(auto_depth), 'b', lw=0.24, label='cal')

    # 设置图例样式和位置
    set_legend(plt)

    # ========== 与自动拾取初至的差值 ==========

    ax_diff = fig_vel.add_axes([0.47, 0.3, 0.2, 0.55])  # 拉平, left，bottom，width，height

    # 手动拾取
    if param["path_hand_op"]:
        pass_hand = 0
        diff_hand = []
        for i in range(len(auto_depth)):
            diff = hand_time[i] - auto_time[i]
            diff_hand.append(diff)
            if abs(diff) <= 0.5:
                pass_hand = pass_hand + 1
        ax_diff.plot(diff_hand, auto_depth, 'ro-', lw=0.24, markersize=0.3,
                     label=f'hand, {round(float(pass_hand / len(auto_depth)), 4)}')

    # 计算初至
    diff_cal = []
    pass_cal = 0
    for i in range(len(auto_depth)):
        diff = cal_time[i] - auto_time[i]
        diff_cal.append(diff)
        if abs(diff) <= 0.5:
            pass_cal = pass_cal + 1
    ax_diff.plot(diff_cal, auto_depth, 'bo-', lw=0.24, markersize=0.3,
                 label=f'cal, {round(float(pass_cal / len(auto_depth)), 4)}')

    ax_diff.plot([-0.5, -0.5], [0, np.max(auto_depth) + 20], 'g', lw=0.24)
    ax_diff.plot([0.5, 0.5], [0, np.max(auto_depth) + 20], 'g', lw=0.24)
    ax_diff.plot([0, 0], [0, np.max(auto_depth) + 20], 'g', lw=0.24)

    x_min_diff = -1
    x_max_diff = 1
    y_min_diff = np.min(auto_depth) - 20
    y_max_diff = np.max(auto_depth) + 20
    set_axis(ax_diff, fig_vel, plt, x_min_diff, x_max_diff, y_min_diff, y_max_diff, "Diff(ms)", "Deep(m)")

    # 设置图例样式和位置
    set_legend(plt)

    # ========== 与上一个初至的差值 ==========

    last_ax_diff = fig_vel.add_axes([0.75, 0.3, 0.2, 0.55])  # 拉平, left，bottom，width，height

    # 自动拾取
    last_diff_auto = []
    for i in range(len(auto_depth)):
        if i == 0:
            last_diff_auto.append(0)
        else:
            last_diff_auto.append(auto_time[i] - auto_time[i - 1])
    last_diff_auto[0] = last_diff_auto[1]
    last_ax_diff.plot(last_diff_auto, auto_depth, 'g', lw=0.24, markersize=0.3, label='auto')

    # 手动拾取
    if param["path_hand_op"]:
        last_diff_hand = []
        for i in range(len(auto_depth)):
            if i == 0:
                last_diff_hand.append(0)
            else:
                last_diff_hand.append(hand_time[i] - hand_time[i - 1])
        last_diff_hand[0] = last_diff_hand[1]
        last_ax_diff.plot(last_diff_hand, auto_depth, 'r', lw=0.24, markersize=0.3, label='hand')

    # 计算初至
    last_diff_cal = []
    for i in range(len(auto_depth)):
        if i == 0:
            last_diff_cal.append(0)
        else:
            last_diff_cal.append(cal_time[i] - cal_time[i - 1])
    last_diff_cal[0] = last_diff_cal[1]
    last_ax_diff.plot(last_diff_cal, auto_depth, 'b', lw=0.24, markersize=0.3, label='cal')

    # 深度差
    last_ax_diff.plot(np.array(depth_diff) / 10, auto_depth, 'y', lw=0.24, markersize=0.3, label='deep')

    x_min_last_diff = 0
    x_max_last_diff = np.max(last_diff_auto)
    y_min_last_diff = np.min(auto_depth) - 20
    y_max_last_diff = np.max(auto_depth) + 20
    set_axis(last_ax_diff, fig_vel, plt, x_min_last_diff, x_max_last_diff, y_min_last_diff, y_max_last_diff,
             "LastDiff(ms)", "Deep(m)")
    last_ax_diff.invert_xaxis()  # 颠倒x轴

    # 设置图例样式和位置
    set_legend(plt)

    # ========== 深度差分界线 ==========
    for i in range(len(depth_flag)):
        ax_vel.axhline(y=depth_flag[i], color='red', linewidth=0.24, zorder=3)
        ax_diff.axhline(y=depth_flag[i], color='red', linewidth=0.24, zorder=2)
        last_ax_diff.axhline(y=depth_flag[i], color='red', linewidth=0.24, zorder=1)

    plt.close(0)  # 关闭空白窗口
    plt.show()

    if param["out_path"]:
        auto_op['time'] = [float(time) for time in cal_time]
        # auto_op['time'] = cal_time
        save_txt(auto_op, param["out_path"])
        print("==========*****==========\n数据保存完成")

if __name__ == "__main__":
    main({
        # sgy图像
        "path_sgy": r"D:\myFile\appDev\fb2\vsp_data\xy101_z.sgy",
        # "path_sgy": r"",
        "sgy_depth_start": 0,
        "sgy_depth_end": 10000,
        "sgy_time_start": 0,
        "sgy_time_end": 3500,
        # 自动拾取
        "path_auto_op": r"D:\myFile\appDev\fb2\vsp_data\xy101_auto_op_5.23.txt",
        "show_sgy_auto": True,  # 是否在sgy图像中绘制自动拾取的初至
        # 手动拾取
        # "path_hand_op": r"D:\myFile\appDev\fb2\vsp_data\xy101_my_op_5.23.txt",
        "path_hand_op": r"",
        "show_sgy_hand": True,  # 是否在sgy图像中绘制手动拾取的初至
        # 计算初至
        "show_sgy_cal": True,  # 是否在sgy图像中绘制计算得到的初至
        "cal_depth_start": 40,
        "cal_depth_end": 10000,
        # 输出路径
        # "out_path": r"D:\myFile\appDev\fb2\vsp_data\xy101_cal_op_5.26_2.txt",
        "out_path": r""
    })
